import pandas as pd
import matplotlib.pyplot as plt
import os
import numpy as np

# ======================
# 配置
# ======================
log_csv = './logs/monitor.csv'  # Monitor 日志路径
output_dir = './logs/figures'   # 输出图像目录
os.makedirs(output_dir, exist_ok=True)

# ======================
# 读取日志
# ======================
df = pd.read_csv(log_csv, skiprows=1)

# 处理 Monitor 文件中列名可能带空格或首尾特殊字符
df.columns = [c.strip() for c in df.columns]

# 步数累计
steps = df['l'].cumsum()  # 'l' 是每个 episode 的长度

# ======================
# 奖励变化
# ======================
plt.figure(figsize=(12,5))
plt.plot(steps, df['r'], label='Episode Reward')
plt.xlabel('Steps')
plt.ylabel('Reward')
plt.title('Trading PPO Training Reward')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.savefig(os.path.join(output_dir, "reward_curve.png"))
plt.close()

# ======================
# 完成率变化
# ======================
if 'completion_rate' in df.columns:
    plt.figure(figsize=(12,5))
    plt.plot(steps, df['completion_rate'], label='Completion Rate')
    plt.xlabel('Steps')
    plt.ylabel('Completion Rate')
    plt.title('Trading PPO Completion Rate')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "completion_curve.png"))
    plt.close()

# ======================
# VWAP / 平均成交价变化
# ======================
if 'avg_price' in df.columns:
    plt.figure(figsize=(12,5))
    plt.plot(steps, df['avg_price'], label='Avg Price / VWAP')
    if 'market_price' in df.columns:
        plt.plot(steps, df['market_price'], label='Market Price', linestyle='--')
    plt.xlabel('Steps')
    plt.ylabel('Price')
    plt.title('VWAP / Avg Price During Training')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "vwap_curve.png"))
    plt.close()

# ======================
# 成本偏差变化
# ======================
if 'avg_price' in df.columns and 'market_price' in df.columns:
    cost_deviation = np.abs(df['avg_price'] - df['market_price']) / df['market_price']
    plt.figure(figsize=(12,5))
    plt.plot(steps, cost_deviation, label='Cost Deviation')
    plt.xlabel('Steps')
    plt.ylabel('Deviation')
    plt.title('VWAP vs Market Price Deviation')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "cost_deviation.png"))
    plt.close()

# ======================
# 步数效率变化
# ======================
if 'executed_quantity' in df.columns:
    step_efficiency = df['executed_quantity'] / df['l']
    plt.figure(figsize=(12,5))
    plt.plot(steps, step_efficiency, label='Step Efficiency')
    plt.xlabel('Steps')
    plt.ylabel('Executed Quantity per Step')
    plt.title('Step Efficiency During Training')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "step_efficiency.png"))
    plt.close()

print(f"Visualization finished! Figures are saved to {output_dir}")
